import os
import asyncio
from contextlib import asynccontextmanager
from typing import List, Optional, Dict
from app.core.config import settings

class APIKeyPool:
    def __init__(self, keys: List[str], per_key_slots: int, global_max_concurrency: Optional[int] = None):
        self._keys = [k.strip() for k in keys if k and k.strip()]
        if not self._keys:
            fallback = os.getenv("OPENAI_API_KEY", "").strip()
            if fallback:
                self._keys = [fallback]
        if not self._keys:
            raise ValueError("Zeakai API keys are not configured")

        self._queue: asyncio.Queue[str] = asyncio.Queue()
        for key in self._keys:
            for _ in range(max(1, per_key_slots)):
                self._queue.put_nowait(key)

        self._global_sem = asyncio.Semaphore(global_max_concurrency) if (global_max_concurrency and global_max_concurrency > 0) else None

    @asynccontextmanager
    async def acquire(self):
        if self._global_sem is not None:
            await self._global_sem.acquire()
        key = await self._queue.get()
        try:
            yield key
        finally:
            self._queue.put_nowait(key)
            if self._global_sem is not None:
                self._global_sem.release()

# 用每个事件循环一个池，避免跨循环引用
_loop_pools: Dict[int, APIKeyPool] = {}

def get_zeakai_key_pool() -> APIKeyPool:
    # 保证在当前运行的事件循环里创建
    try:
        loop = asyncio.get_running_loop()
    except RuntimeError:
        loop = asyncio.get_event_loop()
    loop_id = id(loop)

    pool = _loop_pools.get(loop_id)
    if pool is None:
        raw = settings.zeakai_api_keys
        keys = [s.strip() for s in raw.split(",")] if raw else []
        pool = APIKeyPool(
            keys=keys,
            per_key_slots=max(1, settings.zeakai_per_key_max_concurrency),
            global_max_concurrency=settings.zeakai_global_max_concurrency,
        )
        _loop_pools[loop_id] = pool
    return pool